// 导入图像处理相关模块
use image::imageops;
// 导入路径处理模块
use std::path::Path;
// 导入文件系统操作模块
use std::fs;

// 加载水印图片
fn load_watermark(path: &str) -> image::DynamicImage {
    let path_buf = Path::new(path).to_path_buf();
    image::open(&path_buf).expect(("读取水印图片失败".to_string() + path).as_str())
}

// 调整水印尺寸
fn resize_watermark(watermark: &image::DynamicImage) -> image::DynamicImage {
    let width = watermark.width();
    let height = watermark.height();
    let resized = image::imageops::resize(
        watermark,
        width,
        height,
        image::imageops::FilterType::Lanczos3,
    );
    image::DynamicImage::ImageRgba8(resized)
}

// 应用透明度
fn apply_transparency(watermark: &image::DynamicImage, alpha: f32) -> image::DynamicImage {
    let transparent = image::ImageBuffer::from_fn(watermark.width(), watermark.height(), |x, y| {
        let pixel = watermark
            .as_rgba8()
            .expect("无法将图片转换为 RGBA8 格式")
            .get_pixel(x, y);
        image::Rgba([
            pixel[0],
            pixel[1],
            pixel[2],
            (pixel[3] as f32 * alpha) as u8,
        ])
    });
    image::DynamicImage::ImageRgba8(transparent)
}

// 应用水印到图片
fn apply_watermark(img: &mut image::DynamicImage, watermark: &image::DynamicImage, x: i64, y: i64) {
    imageops::overlay(img, watermark, x, y);
}

// 保存带水印的图片
fn save_marked_image(img: &image::DynamicImage, marked_dir: &str, original_path: &std::path::Path) {
    fs::create_dir_all(marked_dir).expect("Failed to create marked directory");
    let marked_path = Path::new(marked_dir).join(
        original_path
            .file_name()
            .unwrap()
            .to_str()
            .expect("Invalid UTF-8 filename")
            .replace(".", "_marked."),
    );
    let marked_path_str = marked_path
        .to_str()
        .expect("Failed to convert PathBuf to &str");
    img.save(&marked_path)
        .expect(("保存图片失败：".to_string() + marked_path_str).as_str());
}

// 旋转水印
fn rotate_watermark(watermark: &image::DynamicImage, angle: f32) -> image::DynamicImage {
    use imageproc::geometric_transformations::{rotate_about_center, Interpolation};

    let radians = angle.to_radians();
    let (width, height) = (watermark.width() as f32, watermark.height() as f32);
    let new_width = (width * radians.cos().abs() + height * radians.sin().abs()).ceil() as u32;
    let new_height = (width * radians.sin().abs() + height * radians.cos().abs()).ceil() as u32;

    // 创建足够大的画布来容纳旋转后的图像
    let mut canvas =
        image::ImageBuffer::from_pixel(new_width, new_height, image::Rgba([0, 0, 0, 0]));

    // 将水印放置在画布中央
    let offset_x = (new_width - watermark.width()) / 2;
    let offset_y = (new_height - watermark.height()) / 2;
    image::imageops::overlay(&mut canvas, watermark, offset_x as i64, offset_y as i64);

    let rotated = rotate_about_center(
        &canvas,
        radians,
        Interpolation::Bilinear,
        image::Rgba([0, 0, 0, 0]),
    );
    image::DynamicImage::ImageRgba8(rotated)
}

// 主函数，程序入口
#[tokio::main]
async fn main() {
    // 获取命令行参数
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 4 {
        eprintln!(
            "使用方法: {} <水印图片路径> <待处理图片目录路径> <输出目录路径>，如果路径包含空格，请用双引号包裹",
            args[0]
        );
        eprintln!("示例: {} watermark.png images output", args[0]);
        std::process::exit(1);
    }
    let watermark_path = &args[1];
    let pic_dir = &args[2];
    let output_dir = &args[3];

    // 验证水印路径是否为文件
    let path = Path::new(watermark_path);
    if !path.exists() {
        panic!("水印图片路径不存在: {}", watermark_path);
    }
    if path.is_dir() {
        panic!("水印图片路径不能是文件夹: {}", watermark_path);
    }

    // 验证图片目录路径是否为文件夹
    let pic_path = Path::new(pic_dir).canonicalize().unwrap_or_else(|_| {
        panic!("图片目录路径不存在: {}", pic_dir);
    });
    if !pic_path.is_dir() {
        panic!("图片目录路径必须是文件夹: {}", pic_dir);
    }
    if !pic_path.is_dir() {
        panic!("图片目录路径必须是文件夹: {}", pic_dir);
    }

    // 验证输出目录路径是否为文件夹
    let output_path = Path::new(output_dir);
    if !output_path.exists() {
        fs::create_dir_all(output_dir).expect("创建输出目录失败");
    }
    if !output_path.is_dir() {
        panic!("输出目录路径必须是文件夹: {}", output_dir);
    }

    // 加载水印图片
    let mut watermark = load_watermark(watermark_path);

    // 设置水印参数
    let x = 100;
    let y = 100;
    let alpha = 0.5;
    let rotation_angle = 45.0;

    // 处理水印
    watermark = resize_watermark(&watermark);
    watermark = apply_transparency(&watermark, alpha);
    watermark = rotate_watermark(&watermark, rotation_angle);

    // 使用从命令行参数获取的图片目录路径
    // 收集所有图片路径
    let entries: Vec<_> = fs::read_dir(pic_dir).expect("读取图片目录失败").collect();
    // 使用Tokio多线程并行处理
    let watermark = std::sync::Arc::new(watermark);

    let mut handles = Vec::new();
    for entry in entries {
        let entry = entry.expect("读取图片目录失败");
        let path = entry.path();
        if path.is_file() {
            let watermark = watermark.clone();
            let x = x;
            let y = y;
            let marked_dir = output_dir.to_string();
            handles.push(tokio::task::spawn_blocking(move || {
                // 打开图片
                let marked_dir_path = Path::new(&marked_dir).to_path_buf();
                // 检查文件是否为图片类型
                let ext = match path.extension() {
                    Some(ext) => ext.to_str().unwrap_or(""),
                    None => "",
                };
                if !["png", "jpg", "jpeg", "gif"].contains(&ext.to_lowercase().as_str()) {
                    println!("跳过非图片文件: {}", path.display());
                    return;
                }
                let mut img = image::open(&path).expect(
                    ("读取要加水印的图片失败：".to_string() + path.to_str().unwrap()).as_str(),
                );

                apply_watermark(&mut img, &*watermark, x, y);
                save_marked_image(
                    &img,
                    marked_dir_path
                        .to_str()
                        .expect("Failed to convert PathBuf to &str"),
                    &path,
                );
                println!("{} 加水印完成", path.file_name().unwrap().to_str().unwrap());
            }));
        }
    }
    // 等待所有任务完成
    for handle in handles {
        handle.await.expect("Task execution failed");
    }
    println!("所有图片处理完成");
}
